/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include <iostream>
#include <vector>
#include "data_utils.h"
#include "acl/acl.h"
#include "atb/context.h"
#include "atb/types.h"
#include "atb/svector.h"
#include "atb/operation.h"
#include "atb/infer_op_params.h"
#include "../../../examples/atb_plugin/plugin_custom_operations/layernorm_operation.h"
#include "../../../examples/atb_plugin/plugin_aclnn_operations/argmax_operation.h"
#include "../../../examples/atb_plugin/plugin_aclnn_operations/cast_operation.h"

const float LAYERNORM_FACTOR = 0.0009765625;
const float LAYERNORM_MFACTOR = -0.0009765625;
const float LAYERNORM_EPSILON = 1e-5;

int64_t GetShapeSize(const std::vector<int64_t>& shape)
{
    int64_t shapeSize = 1;
    for (auto i : shape) {
        shapeSize *= i;
    }
    return shapeSize;
}

atb::Tensor CreateAtbTensor(uint8_t* deviceAddr, uint8_t* hostAddr, const std::vector<int64_t>& shape,
                            aclDataType dataType, aclFormat dataFormat)
{
    atb::Tensor tensor;
    tensor.desc.dtype = dataType;
    tensor.desc.format = dataFormat;
    for (uint64_t i = 0; i < shape.size(); i++) {
        tensor.desc.shape.dims[i] = shape[i];
        tensor.desc.shape.dimNum ++;
    }
    tensor.deviceData = deviceAddr;
    tensor.hostData = hostAddr;
    tensor.dataSize = GetShapeSize(shape) * GetElemSizeByDtype(dataType);
    return tensor;
}

atb::Tensor PrepareInput(const std::vector<int64_t>& shape, const std::string &filePath,
                         aclDataType dataType, aclrtStream stream)
{
    size_t dataSize = GetElemSizeByDtype(dataType);
    void* inputHost = nullptr;
    aclrtMallocHost((void**)(&inputHost), GetShapeSize(shape) * dataSize);
    ReadFile(filePath, inputHost, GetShapeSize(shape) * dataSize);
    uint8_t* inputAddr = nullptr;
    aclrtMalloc((void**)&inputAddr, GetShapeSize(shape) * dataSize, ACL_MEM_MALLOC_HUGE_FIRST);
    aclrtMemcpyAsync(inputAddr, GetShapeSize(shape) * dataSize, inputHost,
                     GetShapeSize(shape) * dataSize, ACL_MEMCPY_HOST_TO_DEVICE, stream);
    atb::Tensor tensor = CreateAtbTensor(inputAddr, nullptr, shape, dataType, aclFormat::ACL_FORMAT_ND);
    aclrtSynchronizeStream(stream);
    aclrtFreeHost(inputHost);
    return tensor;
}

atb::Tensor PrepareOutput(const std::vector<int64_t>& shape, aclDataType dataType)
{
    uint8_t* outputAddr = nullptr;
    aclrtMalloc((void**)&outputAddr, GetShapeSize(shape) * sizeof(float), ACL_MEM_MALLOC_HUGE_FIRST);
    atb::Tensor tensor = CreateAtbTensor(outputAddr, nullptr, shape, dataType, aclFormat::ACL_FORMAT_ND);
    return tensor;
}

int Init(int32_t deviceId, aclrtContext &contextAclrt, atb::Context* &contextAtb, aclrtStream &stream)
{
    aclInit(nullptr);
    aclrtSetDevice(deviceId);
    aclrtCreateContext(&contextAclrt, deviceId);
    atb::CreateContext(&contextAtb);
    aclrtSetCurrentContext(contextAclrt);
    aclrtCreateStream(&stream);
    contextAtb->SetExecuteStream(stream);
    return 0;
}

int Finalize(int32_t deviceId, aclrtContext contextAclrt, atb::Context* contextAtb, aclrtStream stream)
{
    aclrtDestroyStream(stream);
    aclrtDestroyContext(contextAclrt);
    DestroyContext(contextAtb);
    aclrtResetDevice(deviceId);
    aclFinalize();
    return 0;
}

void CreatePAPAAGraphOperation(atb::GraphParam &graphParam, atb::Operation **opGraph)
{
    // 构建Graph PluginOp + ATBOp + PluginOp + ATBOp + ATBOp
    graphParam.inTensorNum = 4;
    graphParam.outTensorNum = 1;
    graphParam.internalTensorNum = 4;
    graphParam.nodes.resize(5);

    size_t nodeId = 0;
    atb::Node &layerNormNode = graphParam.nodes.at(nodeId++);
    atb::Node &addNode = graphParam.nodes.at(nodeId++);
    atb::Node &argMaxNode = graphParam.nodes.at(nodeId++);
    atb::Node &castNode = graphParam.nodes.at(nodeId++);
    atb::Node &castNode2 = graphParam.nodes.at(nodeId++);

    atb_plugin::LayernormParam layerNormParam;
    layerNormParam.factor = LAYERNORM_FACTOR;
    layerNormParam.mfactor = LAYERNORM_MFACTOR;
    layerNormParam.eps = LAYERNORM_EPSILON;
    atb::Operation *LayernormOp = new atb_plugin::LayernormOperation("Layernorm", layerNormParam);
    layerNormNode.operation = LayernormOp;
    layerNormNode.inTensorIds = {0, 1, 2};
    layerNormNode.outTensorIds = {5};

    atb::infer::ElewiseParam addParam;
    addParam.elewiseType = atb::infer::ElewiseParam::ElewiseType::ELEWISE_ADD;
    atb::CreateOperation(addParam, &addNode.operation);
    addNode.inTensorIds = {5, 3};
    addNode.outTensorIds = {6};

    atb_plugin::ArgmaxParam argmaxParam;
    argmaxParam.dim = 1;
    atb::Operation *argmaxOp = new atb_plugin::ArgmaxOperation("Argmax", argmaxParam);
    argMaxNode.operation = argmaxOp;
    argMaxNode.inTensorIds = {6};
    argMaxNode.outTensorIds = {7};

    atb::infer::ElewiseParam castParam;
    castParam.elewiseType = atb::infer::ElewiseParam::ElewiseType::ELEWISE_CAST;
    castParam.outTensorType = aclDataType::ACL_INT64;
    atb::CreateOperation(castParam, &castNode.operation);
    castNode.inTensorIds = {7};
    castNode.outTensorIds = {8};

    atb::infer::ElewiseParam castParam2;
    castParam2.elewiseType = atb::infer::ElewiseParam::ElewiseType::ELEWISE_CAST;
    castParam2.outTensorType = aclDataType::ACL_INT32;
    atb::CreateOperation(castParam2, &castNode2.operation);
    castNode2.inTensorIds = {8};
    castNode2.outTensorIds = {4};

    atb::CreateOperation(graphParam, opGraph);
}

void CreateAPPAAGraphOperation(atb::GraphParam &graphParam, atb::Operation **opGraph)
{
    // 构建Graph ATBOp + PluginOp + PluginOp + ATBOp + ATBOp
    graphParam.inTensorNum = 4;
    graphParam.outTensorNum = 1;
    graphParam.internalTensorNum = 4;
    graphParam.nodes.resize(5);

    size_t nodeId = 0;
    atb::Node &addNode = graphParam.nodes.at(nodeId++);
    atb::Node &layerNormNode = graphParam.nodes.at(nodeId++);
    atb::Node &argMaxNode = graphParam.nodes.at(nodeId++);
    atb::Node &castNode = graphParam.nodes.at(nodeId++);
    atb::Node &castNode2 = graphParam.nodes.at(nodeId++);

    atb::infer::ElewiseParam addParam;
    addParam.elewiseType = atb::infer::ElewiseParam::ElewiseType::ELEWISE_ADD;
    atb::CreateOperation(addParam, &addNode.operation);
    addNode.inTensorIds = {0, 3};
    addNode.outTensorIds = {5};

    atb_plugin::LayernormParam layerNormParam;
    layerNormParam.factor = LAYERNORM_FACTOR;
    layerNormParam.mfactor = LAYERNORM_MFACTOR;
    layerNormParam.eps = LAYERNORM_EPSILON;
    atb::Operation *LayernormOp = new atb_plugin::LayernormOperation("Layernorm", layerNormParam);
    layerNormNode.operation = LayernormOp;
    layerNormNode.inTensorIds = {5, 1, 2};
    layerNormNode.outTensorIds = {6};

    atb_plugin::ArgmaxParam argmaxParam;
    argmaxParam.dim = 1;
    atb::Operation *argmaxOp = new atb_plugin::ArgmaxOperation("Argmax", argmaxParam);
    argMaxNode.operation = argmaxOp;
    argMaxNode.inTensorIds = {6};
    argMaxNode.outTensorIds = {7};

    atb::infer::ElewiseParam castParam;
    castParam.elewiseType = atb::infer::ElewiseParam::ElewiseType::ELEWISE_CAST;
    castParam.outTensorType = aclDataType::ACL_INT64;
    atb::CreateOperation(castParam, &castNode.operation);
    castNode.inTensorIds = {7};
    castNode.outTensorIds = {8};

    atb::infer::ElewiseParam castParam2;
    castParam2.elewiseType = atb::infer::ElewiseParam::ElewiseType::ELEWISE_CAST;
    castParam2.outTensorType = aclDataType::ACL_INT32;
    atb::CreateOperation(castParam2, &castNode2.operation);
    castNode2.inTensorIds = {8};
    castNode2.outTensorIds = {4};

    atb::CreateOperation(graphParam, opGraph);
}

int main()
{
    int32_t deviceId = 0;
    aclrtContext contextAclrt;
    atb::Context* contextAtb;
    aclrtStream stream;
    Init(deviceId, contextAclrt, contextAtb, stream);

    std::vector<int64_t> inputShape = {16384, 1024};
    std::vector<int64_t> paramShape = {1024};
    std::vector<int64_t> outShape = {16384};
    atb::Tensor inputx = PrepareInput(inputShape, "../input/input_x.bin", aclDataType::ACL_FLOAT, stream);
    atb::Tensor inputy = PrepareInput(inputShape, "../input/input_x.bin", aclDataType::ACL_FLOAT, stream);
    atb::Tensor gamma = PrepareInput(paramShape, "../input/input_gamma.bin", aclDataType::ACL_FLOAT, stream);
    atb::Tensor beta = PrepareInput(paramShape, "../input/input_beta.bin", aclDataType::ACL_FLOAT, stream);
    atb::Tensor castOut = PrepareOutput(outShape, aclDataType::ACL_INT32);

    // 构建 PluginOp + ATBOp + PluginOp + ATBOp + ATBOp 图
    atb::Operation *graphPAPAAOp = nullptr;
    atb::GraphParam graphPAPAAParam;
    CreatePAPAAGraphOperation(graphPAPAAParam, &graphPAPAAOp);
    atb::VariantPack v_pack_PAPAA;
    v_pack_PAPAA.inTensors = {inputx, gamma, beta, inputy};
    v_pack_PAPAA.outTensors = {castOut};
    uint8_t* workspaceAddrPAPAA = nullptr;
    uint64_t workspaceSizePAPAA = 0;
    graphPAPAAOp->Setup(v_pack_PAPAA, workspaceSizePAPAA, contextAtb);
    std::cout << "graphPAPAAOp workspace size: " << workspaceSizePAPAA << std::endl;
    if (workspaceSizePAPAA) {
        aclrtMalloc((void**)&workspaceAddrPAPAA, workspaceSizePAPAA, ACL_MEM_MALLOC_HUGE_FIRST);
    }
    graphPAPAAOp->Execute(v_pack_PAPAA, workspaceAddrPAPAA, workspaceSizePAPAA, contextAtb);
    if (workspaceAddrPAPAA != nullptr) {
        aclrtFree(workspaceAddrPAPAA);
    }

    // 构建 ATBOp + PluginOp + PluginOp + ATBOp + ATBOp 图
    // atb::Operation *graphAPPAAOp = nullptr;
    // atb::GraphParam graphAPPAAParam;
    // CreateAPPAAGraphOperation(graphAPPAAParam, &graphAPPAAOp);
    // atb::VariantPack v_pack_APPAA;
    // v_pack_APPAA.inTensors = {inputx, gamma, beta, inputy};
    // v_pack_APPAA.outTensors = {castOut};
    // uint8_t* workspaceAddrAPPAA = nullptr;
    // uint64_t workspaceSizeAPPAA = 0;
    // graphAPPAAOp->Setup(v_pack_APPAA, workspaceSizeAPPAA, contextAtb);
    // std::cout << "graphAPPAAOp workspace size: " << workspaceSizeAPPAA << std::endl;
    // if (workspaceSizeAPPAA) {
    //     aclrtMalloc((void**)&workspaceAddrAPPAA, workspaceSizeAPPAA, ACL_MEM_MALLOC_HUGE_FIRST);
    // }
    // graphAPPAAOp->Execute(v_pack_APPAA, workspaceAddrAPPAA, workspaceSizeAPPAA, contextAtb);
    // if (workspaceAddrAPPAA != nullptr) {
    //     aclrtFree(workspaceAddrAPPAA);
    // }

    aclrtFree(inputx.deviceData);
    aclrtFree(gamma.deviceData);
    aclrtFree(beta.deviceData);
    aclrtFree(castOut.deviceData);
    Finalize(deviceId, contextAclrt, contextAtb, stream);
    return 0;
}